from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType
import numpy as np
import openai
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.llms import OpenAI


# 1. 连接 Milvus 数据库
def connect_to_milvus(host="127.0.0.1", port="19530"):
    connections.connect("default", host=host, port=port)
    print("连接成功！")


# 2. 创建索引函数
def create_index(collection: Collection, field_name="embedding"):
    index_params = {
        "index_type": "IVF_FLAT",  # 常见向量索引类型
        "metric_type": "COSINE",  # 相似度度量方法（可以选择 "L2", "IP", "COSINE" 等）
        "params": {"nlist": 1024},  # 你可以根据数据大小调整 nlist 值
    }
    # 创建索引
    collection.create_index(field_name, index_params)
    print("索引创建成功！")


# 3. 加载集合
def load_collection(collection_name: str):
    collection = Collection(collection_name)
    collection.load()  # 加载集合到内存
    print(f"集合 {collection_name} 加载成功！")
    return collection


# 4. 从 Milvus 搜索向量
def search_milvus(collection: Collection, query_embedding, top_k=5):
    # 设置搜索参数，top_k 表示返回最相似的前 K 个结果
    search_params = {
        "metric_type": "COSINE",  # 向量相似度度量方式
        "params": {"nprobe": 16},  # nprobe 决定了搜索的范围
    }
    # 搜索最相似的向量
    results = collection.search(
        data=[query_embedding],  # 查询的向量
        anns_field="embedding",  # 向量字段
        param=search_params,  # 搜索参数
        limit=top_k,  # 返回的数量
        expr=None  # 无需过滤条件
    )
    return results


# 5. 使用 GPT-2 或其他模型生成回答
def generate_answer(query: str, context: str, model="gpt2"):
    # 设置 OpenAI GPT-2 API 的请求
    openai.api_key = "your-openai-api-key"  # 你需要提供 OpenAI 的 API 密钥

    prompt = f"Question: {query}\nContext: {context}\nAnswer:"

    # 请求 OpenAI GPT-2 模型生成回答
    response = openai.Completion.create(
        model=model,
        prompt=prompt,
        max_tokens=150,
        temperature=0.7,
    )
    return response.choices[0].text.strip()


# 6. RAG 查询逻辑
def rag_query(query: str, collection_name: str):
    # 获取 query 的向量表示
    query_embedding = get_query_embedding(query)  # 你可以使用任何模型生成向量

    # 连接到 Milvus 数据库
    connect_to_milvus()

    # 加载集合
    collection = load_collection(collection_name)

    # 搜索与查询最相关的向量
    results = search_milvus(collection, query_embedding)

    # 拼接文档片段作为上下文（用于输入到 GPT-2）
    context = "\n".join([result.entity["text"] for result in results[0]])

    # 使用 GPT-2 生成最终回答
    answer = generate_answer(query, context, model="gpt2")

    return answer


# 模拟查询嵌入
def get_query_embedding(query: str):
    # 你可以使用任何模型生成 query 的嵌入向量。这里使用一个示例
    return np.random.rand(512).tolist()  # 假设嵌入是一个 512 维的向量


# ----------------------
# 主程序：RAG 查询
# ----------------------
if __name__ == "__main__":
    query = "静态资源"  # 用户的查询问题
    collection_name = "doc_chunks"  # 替换为你的集合名称

    answer = rag_query(query, collection_name)
    print("生成的回答：", answer)
